package com.ming.tflitesingleblur.utils;

import android.annotation.SuppressLint;
import android.content.Context;
import android.graphics.Bitmap;
import android.graphics.Canvas;
import android.graphics.Color;
import android.graphics.Paint;
import android.util.DisplayMetrics;
import android.util.Log;
import android.view.WindowManager;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;

/**
 * Created by glue.wang on 2019/12/4.
 * e-mail:miniminiming@qq.com
 */
public class ImageSaveUtils {
    private static String TAG = "ImageSaveUtils";

    /**
     *
     * @param imageData 注意这里的数据源是rgb浮点数组
     * @param shape
     * @param normalized
     * @param normalizeSection 归一化区间如0~1.0扩展到0~255
     * @return
     */
    @SuppressLint("NewApi")
    public static Bitmap convertDataToBitmap(float[] imageData,
                                             int[] shape,
                                             boolean normalized,
                                             float[] normalizeSection) {

        if(normalizeSection == null || normalizeSection.length != 2){
            normalizeSection = new float[]{0, 255f};
        }

        float sectionLow = normalizeSection[0];
        float sectionHigh = normalizeSection[1];

        int[] pixels;
        Bitmap bitmap;
        if (shape[2] == 3) { //rgb三通道
            pixels = new int[imageData.length / 3];
            //把输入数据转化一下
            for (int i = 0; i < imageData.length / 3; i++) {
                int pixel;
                int r;
                int g;
                int b;
                if (normalized) {
                    r = (int) (sectionLow + imageData[i * 3] * sectionHigh);
                    g = (int) (sectionLow + imageData[i * 3 + 1] * sectionHigh);
                    b = (int) (sectionLow + imageData[i * 3 + 2] * sectionHigh);
                } else {
                    r = (int) imageData[i * 3];
                    g = (int) imageData[i * 3 + 1];
                    b = (int) imageData[i * 3 + 2];
                }
                //这里存的是rgb值，需要转化成bitmap所需的像素值，注意参数类型，是int还是float
                pixel = Color.rgb(r, g, b);
                pixels[i] = pixel;
            }
        } else { //这里其实是shape[2] == 1的情况，此时是灰度图，rgb值为同一个
            pixels = new int[imageData.length];
            //把输入数据转化一下
            for (int i = 0; i < imageData.length; i++) {
                int pixel;
                int r, g, b;
                if (normalized) {
                    r = g = b = (int) (sectionLow + imageData[i] * sectionHigh);
                } else {
                    r = g = b = (int) imageData[i];
                }
                //这里存的是rgb值，需要转化成bitmap所需的像素值，注意参数类型，是int还是float
                pixel = Color.rgb(r, g, b);
                pixels[i] = pixel;
            }

        }
        int width = shape[0];
        int height = shape[1];
        bitmap = Bitmap.createBitmap(width, height, Bitmap.Config.RGB_565);
        bitmap.setPixels(pixels, 0, width, 0, 0, width, height);

//        Log.i(TAG, "convertDataToBitmap: imageData " + Arrays.toString(imageData));
//        Log.i(TAG, "convertDataToBitmap: pixels " + Arrays.toString(pixels));


        return bitmap;
    }

    /**
     * 保存图片
     *
     * @param savePath
     * @param imageName
     * @param imageData  数据  直接从TensorImage中取出数组即可
     * @param shape      数据的shape [256,256,3]
     * @param normalized 数据是否有是归一化以后的数据[0.0,1.0]还是[0,255] 前者为true
     */
    public static void saveImage(String savePath,
                                 String imageName,
                                 float[] imageData,
                                 int[] shape,
                                 boolean normalized,
                                 float[] normalizeSection) {
        //如果文件夹不存在，就先创建文件夹
        File dirs=new File(savePath);
        if(!dirs.exists()){//如果文件夹不存在
            dirs.mkdirs();//创建文件夹
        }
        Bitmap bitmap = convertDataToBitmap(imageData, shape, normalized,normalizeSection);

        File imageFile = new File(savePath, imageName);//设置文件名称

        if (imageFile.exists()) {
            imageFile.delete();
        }
        try {
            imageFile.createNewFile();
            FileOutputStream fos = new FileOutputStream(imageFile);
            bitmap.compress(Bitmap.CompressFormat.JPEG, 100, fos);
            fos.flush();
            fos.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

    }

    /**
     * 保存bitmap
     * @param filePath
     * @param bmp
     * @param fileName
     * @return
     */
    public static boolean saveBitmap(Bitmap bmp, String filePath, final String fileName) {
        //如果文件夹不存在，就先创建文件夹
        File dirs=new File(filePath);
        if(!dirs.exists()){//如果文件夹不存在
            dirs.mkdirs();//创建文件夹
        }
        //创建文件并写入
        File f = new File(filePath, fileName);
        if (f.exists()) {
            f.delete();
        }

        Bitmap.CompressFormat format;
        if (fileName.endsWith("png"))
            format = Bitmap.CompressFormat.PNG;
        else
            format = Bitmap.CompressFormat.JPEG;

        try {
            FileOutputStream out = new FileOutputStream(f);
            bmp.compress(format, 80, out);
            out.flush();
            out.close();
            Log.i(TAG, "save complete");

            return true;
        } catch (FileNotFoundException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        } catch (IOException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
        return false;
    }

    public static float[] generateSegmentMask(float[] predictData, int[] shape){
        //这里先默认是3通道的
        int width = shape[0];
        int height = shape[1];
        int channel = shape[2];
        float[] maskResult = new float[width * height];
        //记录每个完整通道里最大数值的那个索引
        for(int i = 0; i < width * height; i++){
            float maxValue = -10000;
            int maxIndex = 0;
            for(int j =0; j< channel; j++){
                if(predictData[i*channel + j] > maxValue){
                    //记录最大数值的参数,这里稍微做一些改动，让0和2都变成0
//                    maxIndex = j == 2 ? 0 : j;
                    maxIndex = j;
                    maxValue = predictData[i*channel + j];
                }
            }
            maskResult[i] = maxIndex;
        }

        return maskResult;
    }

    public static Bitmap convert(Bitmap bitmap, Bitmap.Config config) {
        Bitmap convertedBitmap = Bitmap.createBitmap(bitmap.getWidth(), bitmap.getHeight(), config);
        Canvas canvas = new Canvas(convertedBitmap);
        Paint paint = new Paint();
        paint.setColor(Color.BLACK);
        canvas.drawBitmap(bitmap, 0, 0, paint);
        return convertedBitmap;
    }

    // 获取屏幕宽度
    public static int getDisplayWidth(Context context) {
        DisplayMetrics dm = new DisplayMetrics();
        WindowManager wm = (WindowManager) context.getSystemService(Context.WINDOW_SERVICE);
        wm.getDefaultDisplay().getMetrics(dm);
        return dm.widthPixels;
    }
}
